{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基于采样的路径规划算法——RRT\n",
    "\n",
    "理论知识参考[博客](https://blog.csdn.net/weixin_42301220/article/details/125254061?spm=1001.2014.3001.5501)\n",
    "\n",
    "\n",
    "<center><img src=\"https://img-blog.csdnimg.cn/ba84ca6f60d34b2880ee55f5c54df6a9.png\" width=60%>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import random\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from celluloid import Camera  # 保存动图时用，pip install celluloid\n",
    "%matplotlib qt5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RRT:\n",
    "    \"\"\"\n",
    "    Class for RRT planning\n",
    "    \"\"\"\n",
    "\n",
    "    class Node:\n",
    "        \"\"\"\n",
    "        创建节点\n",
    "        \"\"\"\n",
    "\n",
    "        def __init__(self, x, y):\n",
    "            self.x = x  # 节点坐标\n",
    "            self.y = y\n",
    "            self.path_x = []  # 路径，作为画图的数据\n",
    "            self.path_y = []\n",
    "            self.parent = None  # 父节点\n",
    "\n",
    "    class AreaBounds:\n",
    "        \"\"\"区域大小\n",
    "        \"\"\"\n",
    "\n",
    "        def __init__(self, area):\n",
    "            self.xmin = float(area[0])\n",
    "            self.xmax = float(area[1])\n",
    "            self.ymin = float(area[2])\n",
    "            self.ymax = float(area[3])\n",
    "\n",
    "    def __init__(self,\n",
    "                 start,\n",
    "                 goal,\n",
    "                 obstacle_list,\n",
    "                 rand_area,\n",
    "                 expand_dis=3.0,\n",
    "                 goal_sample_rate=5,\n",
    "                 max_iter=500,\n",
    "                 play_area=None,\n",
    "                 robot_radius=0.0,\n",
    "                 ):\n",
    "        \"\"\"\n",
    "        Setting Parameter\n",
    "\n",
    "        start:起点 [x,y]\n",
    "        goal:目标点 [x,y]\n",
    "        obstacleList:障碍物位置列表 [[x,y,size],...]\n",
    "        rand_area: 采样区域 x,y ∈ [min,max]\n",
    "        play_area: 约束随机树的范围 [xmin,xmax,ymin,ymax]\n",
    "        robot_radius: 机器人半径\n",
    "        expand_dis: 扩展的步长\n",
    "        goal_sample_rate: 采样目标点的概率，百分制.default: 5，即表示5%的概率直接采样目标点\n",
    "\n",
    "        \"\"\"\n",
    "        self.start = self.Node(start[0], start[1])  # 根节点\n",
    "        self.end = self.Node(goal[0], goal[1])\n",
    "        self.min_rand = rand_area[0]\n",
    "        self.max_rand = rand_area[1]\n",
    "        if play_area is not None:\n",
    "            self.play_area = self.AreaBounds(play_area)\n",
    "        else:\n",
    "            self.play_area = None\n",
    "        self.expand_dis = expand_dis\n",
    "        self.goal_sample_rate = goal_sample_rate\n",
    "        self.max_iter = max_iter\n",
    "        self.obstacle_list = obstacle_list\n",
    "        self.node_list = []\n",
    "        self.robot_radius = robot_radius\n",
    "\n",
    "    def planning(self, animation=True, camara=None):\n",
    "        \"\"\"\n",
    "        rrt path planning\n",
    "\n",
    "        animation: flag for animation on or off\n",
    "\n",
    "\n",
    "        camara: 是否保存动图\n",
    "        \"\"\"\n",
    "\n",
    "        # 将起点作为根节点x_{init}​，加入到随机树的节点集合中。\n",
    "        self.node_list = [self.start]\n",
    "        for i in range(self.max_iter):\n",
    "            # 从可行区域内随机选取一个节点x_{rand}\n",
    "            rnd_node = self.sample_free()\n",
    "\n",
    "            # 已生成的树中利用欧氏距离判断距离x_{rand}​最近的点x_{near}。\n",
    "            nearest_ind = self.get_nearest_node_index(self.node_list, rnd_node)\n",
    "            nearest_node = self.node_list[nearest_ind]\n",
    "\n",
    "            # 从x_{near}与x_{rand}的连线方向上扩展固定步长u，得到新节点 x_{new}\n",
    "            new_node = self.steer(nearest_node, rnd_node, self.expand_dis)\n",
    "\n",
    "            # 如果在可行区域内，且x_{near}与x_{new}之间无障碍物\n",
    "            if self.is_inside_play_area(new_node, self.play_area) and \\\n",
    "               self.obstacle_free(new_node, self.obstacle_list, self.robot_radius):\n",
    "                self.node_list.append(new_node)\n",
    "\n",
    "            # 如果此时得到的节点x_new到目标点的距离小于扩展步长，则直接将目标点作为x_rand。\n",
    "            if self.calc_dist_to_goal(self.node_list[-1].x, self.node_list[-1].y) <= self.expand_dis:\n",
    "                final_node = self.steer(\n",
    "                    self.node_list[-1], self.end, self.expand_dis)\n",
    "                if self.obstacle_free(final_node, self.obstacle_list, self.robot_radius):\n",
    "                    # 返回最终路径\n",
    "                    return self.generate_final_course(len(self.node_list) - 1)\n",
    "\n",
    "            if animation and i % 5 == 0:\n",
    "                self.draw_graph(rnd_node, camara)\n",
    "\n",
    "        return None  # cannot find path\n",
    "\n",
    "    def steer(self, from_node, to_node, extend_length=float(\"inf\")):\n",
    "        \"\"\"连线方向扩展固定步长查找x_new\n",
    "\n",
    "        Args:\n",
    "            from_node (_type_): x_near\n",
    "            to_node (_type_): x_rand\n",
    "            extend_length (_type_, optional): 扩展步长u. Defaults to float(\"inf\").\n",
    "\n",
    "        Returns:\n",
    "            _type_: _description_\n",
    "        \"\"\"\n",
    "        # 利用反正切计算角度, 然后利用角度和步长计算新坐标\n",
    "        d, theta = self.calc_distance_and_angle(from_node, to_node)\n",
    "\n",
    "        # 如果$x_{near}$与$x_{rand}$间的距离小于步长，则直接将$x_{rand}$作为新节点$x_{new}$\n",
    "        if extend_length >= d:\n",
    "            new_x = to_node.x\n",
    "            new_y = to_node.y\n",
    "        else:\n",
    "            new_x = from_node.x+math.cos(theta)*extend_length\n",
    "            new_y = from_node.y+math.sin(theta)*extend_length\n",
    "        new_node = self.Node(new_x, new_y)\n",
    "        new_node.path_x = [from_node.x]\n",
    "        new_node.path_y = [from_node.y]\n",
    "        new_node.path_x.append(new_x)\n",
    "        new_node.path_y.append(new_y)\n",
    "\n",
    "        new_node.parent = from_node\n",
    "\n",
    "        return new_node\n",
    "\n",
    "    def generate_final_course(self, goal_ind):\n",
    "        \"\"\"生成路径\n",
    "        Args:\n",
    "            goal_ind (_type_): 目标点索引\n",
    "\n",
    "        Returns:\n",
    "            _type_: _description_\n",
    "        \"\"\"\n",
    "        path = [[self.end.x, self.end.y]]\n",
    "        node = self.node_list[goal_ind]\n",
    "        while node.parent is not None:\n",
    "            path.append([node.x, node.y])\n",
    "            node = node.parent\n",
    "        path.append([node.x, node.y])\n",
    "\n",
    "        return path\n",
    "\n",
    "    def calc_dist_to_goal(self, x, y):\n",
    "        \"\"\"计算(x,y)离目标点的距离\n",
    "        \"\"\"\n",
    "        dx = x - self.end.x\n",
    "        dy = y - self.end.y\n",
    "        return math.hypot(dx, dy)\n",
    "\n",
    "    def sample_free(self):\n",
    "        # 以（100-goal_sample_rate）%的概率随机生长，(goal_sample_rate)%的概率朝向目标点生长\n",
    "        if random.randint(0, 100) > self.goal_sample_rate:\n",
    "            rnd = self.Node(\n",
    "                random.uniform(self.min_rand, self.max_rand),\n",
    "                random.uniform(self.min_rand, self.max_rand))\n",
    "        else:  # goal point sampling\n",
    "            rnd = self.Node(self.end.x, self.end.y)\n",
    "        return rnd\n",
    "\n",
    "    def draw_graph(self, rnd=None, camera=None):\n",
    "        if camera == None:\n",
    "            plt.clf()\n",
    "        # for stopping simulation with the esc key.\n",
    "        plt.gcf().canvas.mpl_connect(\n",
    "            'key_release_event',\n",
    "            lambda event: [exit(0) if event.key == 'escape' else None])\n",
    "        # 画随机点\n",
    "        if rnd is not None:\n",
    "            plt.plot(rnd.x, rnd.y, \"^k\")\n",
    "            if self.robot_radius > 0.0:\n",
    "                self.plot_circle(rnd.x, rnd.y, self.robot_radius, '-r')\n",
    "        # 画已生成的树\n",
    "        for node in self.node_list:\n",
    "            if node.parent:\n",
    "                plt.plot(node.path_x, node.path_y, \"-g\")\n",
    "\n",
    "        # 画障碍物\n",
    "        for (ox, oy, size) in self.obstacle_list:\n",
    "            self.plot_circle(ox, oy, size)\n",
    "\n",
    "        # 如果约定了可行区域，则画出可行区域\n",
    "        if self.play_area is not None:\n",
    "            plt.plot([self.play_area.xmin, self.play_area.xmax,\n",
    "                      self.play_area.xmax, self.play_area.xmin,\n",
    "                      self.play_area.xmin],\n",
    "                     [self.play_area.ymin, self.play_area.ymin,\n",
    "                      self.play_area.ymax, self.play_area.ymax,\n",
    "                      self.play_area.ymin],\n",
    "                     \"-k\")\n",
    "\n",
    "        # 画出起点和目标点\n",
    "        plt.plot(self.start.x, self.start.y, \"xr\")\n",
    "        plt.plot(self.end.x, self.end.y, \"xr\")\n",
    "        plt.axis(\"equal\")\n",
    "        plt.axis([-2, 15, -2, 15])\n",
    "        plt.grid(True)\n",
    "        plt.pause(0.01)\n",
    "        if camera != None:\n",
    "            camera.snap()\n",
    "    # 静态方法无需实例化，也可以实例化后调用，静态方法内部不能调用self.的变量\n",
    "\n",
    "    @staticmethod\n",
    "    def plot_circle(x, y, size, color=\"-b\"):  # pragma: no cover\n",
    "        deg = list(range(0, 360, 5))\n",
    "        deg.append(0)\n",
    "        xl = [x + size * math.cos(np.deg2rad(d)) for d in deg]\n",
    "        yl = [y + size * math.sin(np.deg2rad(d)) for d in deg]\n",
    "        plt.plot(xl, yl, color)\n",
    "\n",
    "    @staticmethod\n",
    "    def get_nearest_node_index(node_list, rnd_node):\n",
    "        dlist = [(node.x - rnd_node.x)**2 + (node.y - rnd_node.y)**2\n",
    "                 for node in node_list]\n",
    "        minind = dlist.index(min(dlist))\n",
    "\n",
    "        return minind\n",
    "\n",
    "    @staticmethod\n",
    "    def is_inside_play_area(node, play_area):\n",
    "\n",
    "        if play_area is None:\n",
    "            return True  # no play_area was defined, every pos should be ok\n",
    "\n",
    "        if node.x < play_area.xmin or node.x > play_area.xmax or \\\n",
    "           node.y < play_area.ymin or node.y > play_area.ymax:\n",
    "            return False  # outside - bad\n",
    "        else:\n",
    "            return True  # inside - ok\n",
    "\n",
    "    @staticmethod\n",
    "    def obstacle_free(node, obstacleList, robot_radius):\n",
    "\n",
    "        if node is None:\n",
    "            return False\n",
    "\n",
    "        for (ox, oy, size) in obstacleList:\n",
    "            dx_list = [ox - x for x in node.path_x]\n",
    "            dy_list = [oy - y for y in node.path_y]\n",
    "            d_list = [dx * dx + dy * dy for (dx, dy) in zip(dx_list, dy_list)]\n",
    "\n",
    "            if min(d_list) <= (size+robot_radius)**2:\n",
    "                return False  # collision\n",
    "\n",
    "        return True  # safe\n",
    "\n",
    "    @staticmethod\n",
    "    def calc_distance_and_angle(from_node, to_node):\n",
    "        \"\"\"计算两个节点间的距离和方位角\n",
    "\n",
    "        Args:\n",
    "            from_node (_type_): _description_\n",
    "            to_node (_type_): _description_\n",
    "\n",
    "        Returns:\n",
    "            _type_: _description_\n",
    "        \"\"\"\n",
    "        dx = to_node.x - from_node.x\n",
    "        dy = to_node.y - from_node.y\n",
    "        d = math.hypot(dx, dy)\n",
    "        theta = math.atan2(dy, dx)\n",
    "        return d, theta\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start \n",
      "found path!!\n"
     ]
    }
   ],
   "source": [
    "def main(gx=6.0, gy=10.0):\n",
    "    print(\"start \")\n",
    "    fig = plt.figure(1)\n",
    "\n",
    "    camera = Camera(fig)  # 保存动图时使用\n",
    "    camera = None  # 不保存动图时，camara为None\n",
    "    show_animation = True\n",
    "    # ====Search Path with RRT====\n",
    "    obstacleList = [(5, 5, 1), (3, 6, 2), (3, 8, 2), (3, 10, 2), (7, 5, 2),\n",
    "                    (9, 5, 2), (8, 10, 1)]  # [x, y, radius]\n",
    "    # Set Initial parameters\n",
    "    rrt = RRT(\n",
    "        start=[0, 0],\n",
    "        goal=[gx, gy],\n",
    "        rand_area=[-2, 15],\n",
    "        obstacle_list=obstacleList,\n",
    "        play_area=[-2, 12, 0, 14],\n",
    "        robot_radius=0.8\n",
    "    )\n",
    "    path = rrt.planning(animation=show_animation, camara=camera)\n",
    "\n",
    "    if path is None:\n",
    "        print(\"Cannot find path\")\n",
    "    else:\n",
    "        print(\"found path!!\")\n",
    "\n",
    "        # Draw final path\n",
    "        if show_animation:\n",
    "            rrt.draw_graph(camera=camera)\n",
    "            plt.grid(True)\n",
    "            plt.pause(0.01)\n",
    "            plt.plot([x for (x, y) in path], [y for (x, y) in path], '-r')\n",
    "            if camera != None:\n",
    "                camera.snap()\n",
    "                animation = camera.animate()\n",
    "                animation.save('trajectory.gif')\n",
    "            plt.show()\n",
    "\n",
    "\n",
    "main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.12 ('gobigger')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "0c7484b3574347463e16b31029466871583b0d4e5c4ad861e8848f2d3746b4de"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
